package org.common.utils.thread;
import org.common.utils.thread.exception.ThreadHashPoolShutdownException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 线程哈希池
 * @author zhouzhibing
 * @date 2025/2/21
 */
public class ThreadHashPool {

    private static final Logger log = LoggerFactory.getLogger(ThreadHashPool.class);

    /**
     * 线程池数组数量
     */
    private final ThreadPoolExecutor[] threadPoolArray;

    /**
     * 开始
     */
    private boolean shutdown;

    /**
     * 关闭线程池时是否抛出异常
     */
    private boolean shutdownException;

    public ThreadHashPool(String threadName) {
        this(new DefaultThreadFactory(ThreadHashPool.class.getSimpleName() + "-" + threadName));
    }

    public ThreadHashPool(DefaultThreadFactory factory) {
        this(factory , ThreadTool.CPU_COUNT);
    }

    public ThreadHashPool(DefaultThreadFactory factory , int cpuCount) {
        threadPoolArray = newThreadArray(factory , cpuCount);
    }

    /**
     * 创建一个新的线程池数组
     * @param factory 线程工厂
     * @param threadCount 线程数量
     * @return 新的线程池数组
     */
    public static ThreadPoolExecutor[] newThreadArray(DefaultThreadFactory factory , int threadCount) {
        ThreadPoolExecutor [] threadPoolArray = new ThreadPoolExecutor[threadCount];
        for (int i = 0; i < threadCount; i++) {
            ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(
                    1,
                    1,
                    0L,
                    TimeUnit.MILLISECONDS,
                    factory.newBlockingQueue(),
                    factory);
            threadPoolArray[i] = threadPoolExecutor;
        }
        return threadPoolArray;
    }

    /**
     * 等待全部完成
     */
    public void waitFinish() {
        while(!isFinish()) {
            try {
                Thread.sleep(10);
                if(log.isDebugEnabled()) {
                    log.debug("ThreadHashPool.waitFinish threadHashPool = {} " , this);
                }
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }

    /**
     * @return 是否完成执行（队列中数量为0）
     */
    public boolean isFinish() {
        return getTotalQueueSize() == 0;
    }

    /**
     * @return 获取所有队列中任务数量
     */
    public int getTotalQueueSize() {
        int totalQueueCount = 0;
        for (ThreadPoolExecutor executor : threadPoolArray) {
            totalQueueCount += executor.getQueue().size();
        }
        return totalQueueCount;
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        for (ThreadPoolExecutor executor : threadPoolArray) {
            sb.append("activeCount=").append(executor.getActiveCount()).append(", ");
            sb.append("completedTaskCount=").append(executor.getCompletedTaskCount()).append(", ");
            sb.append("queueSize=").append(executor.getQueue().size());
            sb.append("\n");
        }
        return sb.toString();
    }

    /**
     * 关闭
     */
    public void shutdown() {
        shutdown = true;
        waitFinish();
        for (ThreadPoolExecutor executor : threadPoolArray) {
            executor.close();
            executor.shutdown();
        }
    }

    /**
     * 获取下标获取线程执行器
     * @param id 使用的hash计算的ID
     * @return 线程执行器
     */
    private ThreadPoolExecutor getExecutor(int id) {
        int index = Math.abs(id % threadPoolArray.length);
        return threadPoolArray[index];
    }

    /**
     * 根据ID hash执行内容
     * @param id hash计算的ID
     * @param runnable 执行内容
     * @return 是否成功添加任务
     */
    public void execute(int id , Runnable runnable) {
        if(shutdown) {
            if(shutdownException)
                return;
            else
                throw new ThreadHashPoolShutdownException("ThreadHashPool is shutdown");
        }
        getExecutor(id).execute(runnable);
    }

    /**
     * 根据ID hash执行内容
     * @param id hash计算的ID
     * @param callable 执行内容
     * @return 处理对象
     */
    public <T> Future<T> submit(int id , Callable<T> callable) {
        if(shutdown) {
            if(shutdownException)
                return null;
            else
                throw new ThreadHashPoolShutdownException("ThreadHashPool is shutdown");
        }
        return getExecutor(id).submit(callable);
    }

    public void closeShutdownException() {
        shutdownException = true;
    }
    public void openShutdownException() {
        shutdownException = false;
    }
}
